package ping

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"math/rand"
	"net"
	"time"

	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv6"
)

const (
	icmpEchoRequest = 8
)

func PingHost(ip string, count int, timeout time.Duration) (*PingResult, error) {
	ipAddr, err := net.ResolveIPAddr("ip", ip)
	if err != nil {
		fmt.Printf("ERROR0: %s\n", err.Error())
		return nil, err
	}

	var conn *net.IPConn
	if ipAddr.IP.To4() != nil {
		conn, err = net.DialIP("ip4:icmp", nil, ipAddr)
	} else {
		conn, err = net.DialIP("ip6:ipv6-icmp", nil, ipAddr)
	}
	if err != nil {
		fmt.Printf("ERROR1: %s\n", err.Error())
		return nil, err
	}
	defer conn.Close()

	results := &PingResult{MinRTT: time.Duration(1<<63 - 1)} // 初始化结果

	for i := 0; i < count; i++ {
		var msg []byte
		if ipAddr.IP.To4() != nil {
			msg = CreateICMPMessage(i)
		} else {
			msg, err = CreateICMPMessageIPv6(i)
			if err != nil {
				fmt.Printf("ERROR2: %s\n", err.Error())
				return nil, err
			}
		}

		start := time.Now() // 记录开始时间
		if _, err := conn.Write(msg); err != nil {
			fmt.Printf("ERROR3: %s\n", err.Error())
			return nil, err
		}

		reply := make([]byte, 1500)
		conn.SetReadDeadline(time.Now().Add(timeout)) // 设置超时
		n, err := conn.Read(reply)
		if err != nil {
			fmt.Printf("ip: %s ERROR4: %s\n", ip, err.Error())
			results.Timeouts++
			continue
		}

		rtt := time.Since(start)
		if rtt < results.MinRTT {
			results.MinRTT = rtt
		}
		if rtt > results.MaxRTT {
			results.MaxRTT = rtt
		}
		results.TotalRTT += rtt
		results.Count++

		fmt.Printf("%s: bytes=%d time=%v seq=%d\n", ip, n, rtt, i)
	}
	if results.Timeouts != count && results.Count != 0 {
		results.AvgRTT = results.TotalRTT / time.Duration(results.Count)
	} else {
		results.AvgRTT = results.MinRTT
	}
	return results, nil
}

func CreateICMPMessage(seq int) []byte {
	msg := IcmpMessage{
		Type: icmpEchoRequest,
		Code: 0,
		ID:   uint16(rand.Intn(0xffff)),
		Seq:  uint16(seq),
		Data: []byte("PING"),
	}

	msg.Checksum = CalculateChecksum(msg)

	var buffer bytes.Buffer
	binary.Write(&buffer, binary.BigEndian, msg.Type)
	binary.Write(&buffer, binary.BigEndian, msg.Code)
	binary.Write(&buffer, binary.BigEndian, msg.Checksum)
	binary.Write(&buffer, binary.BigEndian, msg.ID)
	binary.Write(&buffer, binary.BigEndian, msg.Seq)
	binary.Write(&buffer, binary.BigEndian, msg.Data)

	return buffer.Bytes()
}

func CreateICMPMessageIPv6(seq int) ([]byte, error) {
	pingMessage := icmp.Message{
		Type: ipv6.ICMPTypeEchoRequest,
		Code: 0,
		Body: &icmp.Echo{
			ID:   rand.Intn(0xffff),
			Seq:  seq,
			Data: []byte("PING"),
		},
	}

	messageBytes, err := pingMessage.Marshal(nil)
	if err != nil {
		return nil, err
	}

	return messageBytes, nil
}

func CalculateChecksum(msg IcmpMessage) uint16 {
	var buffer bytes.Buffer
	binary.Write(&buffer, binary.BigEndian, msg.Type)
	binary.Write(&buffer, binary.BigEndian, msg.Code)
	binary.Write(&buffer, binary.BigEndian, msg.ID)
	binary.Write(&buffer, binary.BigEndian, msg.Seq)
	binary.Write(&buffer, binary.BigEndian, msg.Data)

	data := buffer.Bytes()
	var sum uint32
	for i := 0; i < len(data)-1; i += 2 {
		sum += uint32(data[i])<<8 + uint32(data[i+1])
	}

	if len(data)%2 == 1 {
		sum += uint32(data[len(data)-1]) << 8
	}

	for (sum >> 16) > 0 {
		sum = (sum & 0xffff) + (sum >> 16)
	}

	return uint16(^sum)
}
